import cv2
import numpy as np
from ultralytics import YOLO

# 加载YOLOv11模型
model = YOLO("yolo11n.pt")

def yolo_detect(image): # 检测函数
    results = model(image, verbose=False) # 模型预测
    light_color = None # 默认无颜色
    for result in results: # 遍历结果
        cls = result.boxes.cls.tolist() # 获取类别
        if len(cls) > 0: # 如果有结果
            cls = cls[0] # 获取类别
            if cls == 9: # 如果是交通灯
                # 获取置信度
                conf = result.boxes.conf.tolist()[0] # 获取置信度
                xywh = result.boxes.xywh.tolist()[0] # 获取坐标
                if conf > 0.5: # 如果置信度大于0.5
                    light_color = detect_light_color(image.copy(), xywh) # 识别颜色
    return light_color, False  # 移除 person_Warning 逻辑，默认返回 False

# 识别红绿灯颜色
def detect_light_color(image, xywh):
    """
        传入识别到的结果
    :param image: 传入原始图片 用来进行roi裁剪得出交通灯区域
    :param xywh: 交通灯所在区域的中心点坐标 以及宽度和高度
    :return: light_color 0 红色 1 黄色 2 绿色
    """
    print('traffic light')

    # 通过中心点坐标 取一个红绿灯图
    top_left = (int(xywh[0] - xywh[2] // 2), int(xywh[1] - xywh[3] // 2))
    bottom_right = (int(xywh[0] + xywh[2] // 2), int(xywh[1] + xywh[3] // 2))
    cv2.rectangle(image, top_left, bottom_right, (255, 0, 0), 2)
    """将红绿灯切出来"""
    light_img = image[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
    # HSV空间转换 将RGB颜色空间下的图像转换为HSV颜色空间下的图像
    hsv_image_np = cv2.cvtColor(light_img, cv2.COLOR_BGR2HSV)
    red1_low = [0, 43, 46]
    red1_up = [10, 255, 255]
    red2_low = [156, 43, 46]
    red2_up = [180, 255, 255]

    yellow_low = [26, 43, 46]
    yellow_up = [34, 255, 255]

    green_low = [60, 50, 70]
    green_up = [70, 255, 255]
    mask1 = cv2.inRange(hsv_image_np, np.array(red1_low), np.array(red1_up))
    mask1 += cv2.inRange(hsv_image_np, np.array(red2_low), np.array(red2_up))
    # 记录红色像素点个数
    red_cnt = np.sum(mask1 == 255)
    print('红色像素点数量为：', red_cnt)
    mask2 = cv2.inRange(hsv_image_np, np.array(yellow_low), np.array(yellow_up))
    yellow_cnt = np.sum(mask2 == 255)
    print('黄色像素点为数量：', yellow_cnt)
    mask3 = cv2.inRange(hsv_image_np, np.array(green_low), np.array(green_up))
    green_cnt = np.sum(mask3 == 255)
    print('绿色像素点为数量：', green_cnt)

    light_color = np.argmax([red_cnt, yellow_cnt, green_cnt])
    return light_color